use apikit::{ApiError, StatusCode};
use axum::{
    extract::{FromRequestParts, ws::rejection::WebSocketUpgradeRejection},
    http::request::Parts,
    response::IntoResponse,
};
use thiserror::Error;

use crate::param::WsUpgrade;

use super::error_to_response;

#[derive(Debug, Error, ApiError)]
pub enum WsRejection {
    #[status_code(StatusCode::UPGRADE_REQUIRED)]
    #[error(transparent)]
    Axum(#[from] WebSocketUpgradeRejection),
}

impl IntoResponse for WsRejection {
    fn into_response(self) -> axum::response::Response {
        error_to_response(self)
    }
}

impl<S, SendMsg, ReceiveMsg, SC, RC> FromRequestParts<S> for WsUpgrade<SendMsg, ReceiveMsg, SC, RC>
where
    S: Send + Sync,
{
    type Rejection = WsRejection;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let inner = axum::extract::ws::WebSocketUpgrade::from_request_parts(parts, state).await?;
        Ok(Self::new(Box::new(inner)))
    }
}
